from flask import Flask, render_template, request, redirect, url_for,send_from_directory,send_file
from flask import session,jsonify
from pydub import AudioSegment
from scipy.io import wavfile

import numpy as np
import threading
import datetime
import requests
import librosa
import random
import shutil
import wave
import json
import time
import os 
import re

app = Flask(__name__)

app.secret_key = ''

with open('config/配置文件/姓.txt','r',encoding = 'utf-8') as f:
    first_names = [word.strip() for word in f.readlines() if word.strip()]

with open('config/配置文件/男名.txt','r',encoding = 'utf-8') as f:
    male_last_names = [word.strip() for word in f.readlines() if word.strip()]

with open('config/配置文件/女名.txt','r',encoding = 'utf-8') as f:
    female_last_names = [word.strip() for word in f.readlines() if word.strip()]

with open('config/配置文件/过滤词.txt','r',encoding = 'utf-8') as f:
    filter_key_words = [word.strip() for word in f.readlines() if word.strip()]

with open('config/配置文件/数字.txt','r',encoding = 'utf-8') as f:
    replace_key_words = [word.strip().split('：') for word in f.readlines() if word.strip()]

with open('config/配置文件/替换词.txt','r',encoding = 'utf-8') as f:
    replace_key_words += [word.strip().split('：') for word in f.readlines() if word.strip()]

with open('config/配置文件/保留词.txt','r',encoding = 'utf-8') as f:
    keep_words = {word.strip():''.join(random.sample('abcdefghijklmnopqrstuvwxyz',5)) for word in f.readlines() if word.strip()}

with open('config/配置文件/报时.txt','r',encoding = 'utf-8') as f:
    baoshi_info = {word.strip().split('：')[0] : word.strip().split('：')[1] for word in f.readlines() if word.strip()}


female_names = [a + b + c for a in first_names for b in female_last_names for c in female_last_names]
male_names = [a + b + c for a in first_names for b in male_last_names for c in male_last_names]
yonghu_names = ['用户' + ''.join(random.sample(['一','二','三','四','五','六','七','八','九'],4)) for i in range(10000)]


# 扩展句子变量
def expand_sentence(s):
    global female_names,male_names,yonghu_names

    s = s.replace('|','|')
    
    n = random.uniform(0,1)

    # {昵称}
    for i in range(10):
        if n > 0.66:
            s = re.sub('{昵称}',random.choice(female_names) + random.choice(['宝宝','姐姐']),s, count = 1)
        elif n > 0.33:
            s = re.sub('{昵称}', random.choice(male_names) + random.choice(['宝宝','哥哥']),s, count = 1)
        else:
            s = re.sub('{昵称}', random.choice(yonghu_names) + random.choice(['宝宝']),s, count = 1)
    
    # {xxxx}
    for item in re.findall('\{.+?\}',s):
        if random.random() > 0.5:
            s = s.replace(item,item.strip('{}'),1)
        else:
            s = s.replace(item,'',1)
    
    # <A、B、C、D、E>随机排列组合
    for item in re.findall('<.+?>',s):
        data = item.strip('<>').split('、')
        random.shuffle(data)
        s = s.replace(item,'、'.join(data),1)
    
    # <A，B，C，D，E>随机排列组合
    for item in re.findall('<，+?>',s):
        data = item.strip('<>').split('，')
        random.shuffle(data)
        s = s.replace(item,'，'.join(data),1)
    
    # [A|B|C]随机抽取一个
    for match in re.findall("\[.+?\]", s):
        s = s.replace(match,random.choice(match.strip('[').strip(']').split('|')),1)
    
    # 【A|B|C】随机抽取一个
    for match in re.findall("【.+?】", s):
        s = s.replace(match,random.choice(match.strip('【').strip('】').split('|')),1)
    
    return s.strip()


# 平台列表
platforms = []

for platform in os.listdir(os.path.join('config','话术脚本')):
    if platform.startswith('.'):
        continue

    platforms.append(platform)

# 产品列表
products = []

for platform in os.listdir(os.path.join('config','话术脚本')):
    if platform.startswith('.'):
        continue

    for product in os.listdir(os.path.join('config','话术脚本',platform)):
        if product.startswith('.'):
            continue

        products.append(product)

# 音色列表
streamers = [streamer for streamer in os.listdir(os.path.join('D:\\','音色定制')) if os.path.isdir(os.path.join('D:\\','音色定制',streamer))]

task_config = {}

for scene in ['scene-1','scene-2','scene-3','scene-4','scene-5','scene-6','scene-7','scene-8','scene-9','scene-10']:
    task_config[scene] = {
        'scene_name': scene,
        'current_product': '',
        'current_streamer': '',
        'live_platform':'',
        'startGenAudio': False,
        'startPreReplyAudio' : False,
        'volume_add':'',
        'tempScript' : [],
        'script_ip_address' : '192.168.1.2',
        'active': True if scene == 'scene-1' else False
    }


# 打印分割线
def print_split_line():
    width = os.get_terminal_size().columns
    text = '=' * width
    print(text)


def add_volume(input_path,output_path,db_add):
    if db_add == '' or db_add == '0':
        return

    # 加载音频文件
    audio = AudioSegment.from_wav(input_path)

    louder_audio = audio + int(db_add)

    # 保存修改后的音频文件
    louder_audio.export(output_path, format = "wav")


# 清洗话术
def clear_content(content):
    global keep_words,replace_key_words,filter_key_words

    content = content.replace('<think>','')
    content = content.replace('</think>','')
    content = content.replace('“','')
    content = content.lower().strip()
    
    if not content:
        return

    if re.search('[\d\.]+-[\d\.]+',content):
        content = re.sub('-','到',content)
    
    # 保留词
    for k in keep_words:
        content = re.sub(k,keep_words[k],content)
    
    # 替换词
    for rep_word in replace_key_words:
        content = re.sub(rep_word[0].strip(),rep_word[1].strip(),content)

    # 还原保留词
    for k in keep_words:
        content = re.sub(keep_words[k],k,content)

    # 遇到过滤词（涉及到不合逻辑词）和敏感词（涉及到广告法）就直接跳过这一整句
    for fk in filter_key_words:
        if fk in content:
            print(fk,'是过滤词')
            return ''
    
    # 检测到英文
    if re.search('[a-z]{2,}',content):
        return ''

    # 如果替换数字之后，还是检测到数字，那么这一句就先略过
    if re.search('[0-9]',content):
        return ''

    if len(content) > 1000:
        print(content,'太长，不适合生成音频..........')
        return ''

    return content.strip()


# 解析音色克隆文件
def parse_ref_audio_text(scene_idx):
    global task_config

    wav_path_list = [file for file in os.listdir(os.path.join('D:\\','音色定制',task_config[scene_idx]['current_streamer'],'wavs')) if file.endswith('.wav') or file.endswith('.WAV') or file.endswith('.MP3') or file.endswith('.mp3')]

    wav_path_random = random.choice(wav_path_list)

    ref_wav = os.path.abspath(os.path.join('D:\\','音色定制',task_config[scene_idx]['current_streamer'],'wavs',wav_path_random))
    ref_txt = os.path.abspath(os.path.join('D:\\','音色定制',task_config[scene_idx]['current_streamer'],'txts',wav_path_random + '.txt'))

    with open(ref_txt,'r+',encoding = 'utf-8') as f:
        ref_txt = f.read().strip()

    return ref_wav,ref_txt


# 解析音色克隆文件
def parse_ref_audio_text_by_streamer(streamer):
    global task_config

    result = []

    wav_path_list = [file for file in os.listdir(os.path.join('D:\\','音色定制',streamer,'wavs')) if file.endswith('.wav') or file.endswith('.WAV') or file.endswith('.MP3') or file.endswith('.mp3')]

    for wav_path in wav_path_list:
        ref_wav = os.path.abspath(os.path.join('D:\\','音色定制',streamer,'wavs',wav_path))
        ref_txt = os.path.abspath(os.path.join('D:\\','音色定制',streamer,'txts',wav_path + '.txt'))

        with open(ref_txt,'r+',encoding = 'utf-8') as f:
            ref_txt = f.read().strip()

        result.append({
            'ref_wav':ref_wav,
            'ref_txt':ref_txt
        })

    return result


# 读取当前北京时间
def read_current_time():
    global baoshi_info

    formats = [
        "现在是北京时间%Y年%m月%d日%H点%M分，",
        "现在是%Y年%m月%d日%H点%M分，",
        "现在时间来到了%H点%M分，",
        "时间来到了%H点%M分，",
        "现在时间是%H点%M分，"
    ]

    tinfo = datetime.datetime.now().strftime(random.choice(formats))

    for k in baoshi_info:
        if k in tinfo:
            tinfo = tinfo.replace(k,baoshi_info[k])

    return tinfo


# 开始生成主音频
def startGenMainAudioTask(scene_idx):
    global task_config

    for i in range(10000000):
        if len(os.listdir(os.path.join('config','话术音频',task_config[scene_idx]['live_platform'],task_config[scene_idx]['current_product'],task_config[scene_idx]['current_streamer'],'主播'))) > 100:
            print(task_config[scene_idx]['live_platform'],task_config[scene_idx]['current_product'],task_config[scene_idx]['current_streamer'],'剩余音频库存达到100，休眠60秒')
            time.sleep(60)
            continue

        if not task_config[scene_idx]['startGenMainAudioTasking']:
            break

        try:
            print(task_config[scene_idx]['live_platform'],task_config[scene_idx]['current_product'],task_config[scene_idx]['current_streamer'],'请求一个文案')
            print_split_line()

            resp = requests.post('http://' + task_config[scene_idx]['script_ip_address'] + ':6868/startDownloadScript',data = {
                'product': task_config[scene_idx]['current_product'],
                'streamer':task_config[scene_idx]['current_streamer'],
                'platform':task_config[scene_idx]['live_platform']
            }).json()            

        except:
            print(task_config[scene_idx]['live_platform'],task_config[scene_idx]['current_product'],task_config[scene_idx]['current_streamer'],'请求文案失败')
            print_split_line()
            time.sleep(1)
            continue

        if not resp['content'].strip():
            time.sleep(random.uniform(0,1))
            continue

        print(resp['content'])
        
        lines = re.split('[\n| ]',resp['content'])

        print_split_line()

        while True:
            if not lines:
                break

            if not task_config[scene_idx]['startGenMainAudioTasking']:
                break

            gpu_occupied = False

            # 如果其他场景正在生成声音，则表示GPU繁忙
            for scene in task_config:
                if task_config[scene]['startGenAudio']:
                    gpu_occupied = True

            if gpu_occupied:
                time.sleep(random.uniform(0,1))
                continue

            line = lines.pop(0)

            gen_text = clear_content(line) 

            if not gen_text:
                continue

            if len(gen_text) > 280:
                print(gen_text,'字数超过280')
                continue

            # 表示开始生成音频，要占用GPU
            task_config[scene_idx]['startGenAudio'] = True

            ref_wave,ref_text = parse_ref_audio_text(scene_idx)

            print(task_config[scene_idx]['live_platform'],task_config[scene_idx]['current_product'],task_config[scene_idx]['current_streamer'],gen_text)
            print_split_line()

            response = requests.post('http://127.0.0.1:5000/generate_audio',data = {
                'ref_file':ref_wave,
                'ref_text':ref_text,
                'gen_text':gen_text,
                'speed': 1
            })

            # 检查请求是否成功
            if response.status_code == 200:
                save_path = os.path.join('config','话术音频',task_config[scene_idx]['live_platform'],task_config[scene_idx]['current_product'],task_config[scene_idx]['current_streamer'],'主播','%s.wav' % datetime.datetime.now().strftime("%Y_%m_%d_%H_%M_%S"))

                with open(save_path, 'wb') as f4:
                    f4.write(response.content)
                
                time.sleep(1)

                add_volume(save_path,save_path,task_config[scene_idx]['volume_add'])

                # welcome_rate = 0.9

                # if random.uniform(0,1) > welcome_rate:
                #     gen_text = '欢迎' + random.choice(yonghu_names) + '家人，欢迎'  + random.choice(yonghu_names) + '家人，欢迎'  + random.choice(yonghu_names) + '家人，欢迎' + random.choice(yonghu_names) + '家人来到' + random.choice(['我们','工厂','源头','咱们','厂家']) + '直播间。'
                    
                #     print(task_config[scene_idx]['live_platform'],task_config[scene_idx]['current_product'],task_config[scene_idx]['current_streamer'],gen_text)

                #     response = requests.post('http://127.0.0.1:5000/generate_audio',data = {
                #         'ref_file':ref_wave,
                #         'ref_text':ref_text,
                #         'gen_text':gen_text,
                #         'speed': random.uniform(0.9,1.1)
                #     })

                #     # 检查请求是否成功
                #     if response.status_code == 200:
                #         save_path = os.path.join('config','话术音频',task_config[scene_idx]['live_platform'],task_config[scene_idx]['current_product'],task_config[scene_idx]['current_streamer'],'主播','%s.wav' % datetime.datetime.now().strftime("%Y_%m_%d_%H_%M_%S"))

                #         with open(save_path, 'wb') as f5:
                #             f5.write(response.content)

                #         combine_wavs(save_path,'static/5s空白.WAV',save_path)

                #         time.sleep(0.5)
                #         random_audio_volume(save_path,save_path)

                #         time.sleep(0.5)
                #         trim_silence(save_path,save_path)

            # 空出资源让其他任务开始执行
            task_config[scene_idx]['startGenAudio'] = False

            time.sleep(random.uniform(1,2))

    print(task_config[scene_idx]['live_platform'],task_config[scene_idx]['current_streamer'],task_config[scene_idx]['current_product'],'主音频生成结束')

        
@app.route('/')
def login():
    return render_template('login.html')


@app.route('/index')
def index():
    return render_template('index.html', products = products, streamers = streamers,platforms = platforms)


# 开始生成音频
@app.route('/startGenAudio',methods = ['POST'])
def startGenAudio():
    global task_config

    scene_idx = request.form.get('scene_idx')
    script_ip_address = request.form.get('script_ip_address')

    task_config[scene_idx]['script_ip_address'] = script_ip_address

    task_config[scene_idx]['startGenMainAudioTasking'] = True
    task_config[scene_idx]['startGenSubAudioTasking'] = True

    if not os.path.exists(os.path.join('config','话术音频',task_config[scene_idx]['live_platform'])):
        os.mkdir(os.path.join('config','话术音频',task_config[scene_idx]['live_platform']))

    if not os.path.exists(os.path.join('config','话术音频',task_config[scene_idx]['live_platform'],task_config[scene_idx]['current_product'])):
        os.mkdir(os.path.join('config','话术音频',task_config[scene_idx]['live_platform'],task_config[scene_idx]['current_product']))

    if not os.path.exists(os.path.join('config','话术音频',task_config[scene_idx]['live_platform'],task_config[scene_idx]['current_product'],task_config[scene_idx]['current_streamer'])):
        os.mkdir(os.path.join('config','话术音频',task_config[scene_idx]['live_platform'],task_config[scene_idx]['current_product'],task_config[scene_idx]['current_streamer']))

    if not os.path.exists(os.path.join('config','话术音频',task_config[scene_idx]['live_platform'],task_config[scene_idx]['current_product'],task_config[scene_idx]['current_streamer'],'临时')):
        os.mkdir(os.path.join('config','话术音频',task_config[scene_idx]['live_platform'],task_config[scene_idx]['current_product'],task_config[scene_idx]['current_streamer'],'临时'))

    if not os.path.exists(os.path.join('config','话术音频',task_config[scene_idx]['live_platform'],task_config[scene_idx]['current_product'],task_config[scene_idx]['current_streamer'],'主播')):
        os.mkdir(os.path.join('config','话术音频',task_config[scene_idx]['live_platform'],task_config[scene_idx]['current_product'],task_config[scene_idx]['current_streamer'],'主播'))

    # 开始生成主音频
    task_thread = threading.Thread(target = startGenMainAudioTask,args = (scene_idx,))
    task_thread.start()

    return jsonify({'success':True,'message':''})


# 停止生成音频
@app.route('/stopGenAudio',methods = ['POST'])
def stopGenAudio():
    global task_config

    scene_idx = request.form.get('scene_idx')

    task_config[scene_idx]['startGenMainAudioTasking'] = False
    task_config[scene_idx]['startGenSubAudioTasking'] = False
    
    return jsonify({'success':True,'message':''})


# 清空音频
@app.route('/clearAudio',methods = ['POST'])
def clearAudio():
    global task_config

    scene_idx = request.form.get('scene_idx')

    if not os.path.exists(os.path.join('config','话术音频',task_config[scene_idx]['live_platform'])):
        return jsonify({'success':True,'message':''})

    if not os.path.exists(os.path.join('config','话术音频',task_config[scene_idx]['live_platform'],task_config[scene_idx]['current_product'])):
        return jsonify({'success':True,'message':''})

    if not os.path.exists(os.path.join('config','话术音频',task_config[scene_idx]['live_platform'],task_config[scene_idx]['current_product'],task_config[scene_idx]['current_streamer'])):
        return jsonify({'success':True,'message':''})

    if not os.path.exists(os.path.join('config','话术音频',task_config[scene_idx]['live_platform'],task_config[scene_idx]['current_product'],task_config[scene_idx]['current_streamer'],'主播')):
        return jsonify({'success':True,'message':''})

    for file in os.listdir(os.path.join('config','话术音频',task_config[scene_idx]['live_platform'],task_config[scene_idx]['current_product'],task_config[scene_idx]['current_streamer'],'主播')):
        os.remove(os.path.join('config','话术音频',task_config[scene_idx]['live_platform'],task_config[scene_idx]['current_product'],task_config[scene_idx]['current_streamer'],'主播',file))

    for file in os.listdir(os.path.join('config','话术音频',task_config[scene_idx]['live_platform'],task_config[scene_idx]['current_product'],task_config[scene_idx]['current_streamer'],'临时')):
        os.remove(os.path.join('config','话术音频',task_config[scene_idx]['live_platform'],task_config[scene_idx]['current_product'],task_config[scene_idx]['current_streamer'],'临时',file))

    return jsonify({'success':True,'message':''})


# 选择音色
@app.route('/selectStreamer',methods = ['POST'])
def selectStreamer():
    global task_config

    scene_idx = request.form.get('scene_idx')

    current_streamer = request.form.get('streamer')

    task_config[scene_idx]['current_streamer'] = current_streamer

    return jsonify({'success':True,'message':'已切换到' + current_streamer})


# 选择产品
@app.route('/selectProduct',methods = ['POST'])
def selectProduct():
    global task_config

    scene_idx = request.form.get('scene_idx')

    current_product = request.form.get('product')

    task_config[scene_idx]['current_product'] = current_product

    return jsonify({'success':True,'message':'已切换到' + current_product})


# 选择平台
@app.route('/selectLivePlatform',methods = ['POST'])
def selectLivePlatform():
    global task_config

    scene_idx = request.form.get('scene_idx')

    live_platform = request.form.get('live_platform')

    task_config[scene_idx]['live_platform'] = live_platform

    return jsonify({'success':True,'message':'已切换到' + live_platform})


# 选择场景
@app.route('/selectScene',methods = ['POST'])
def selectScene():
    global task_config

    scene_idx = request.form.get('scene_idx')
    
    task_config[scene_idx]['active'] = True

    return jsonify({'success':True,'message':''})


# 音量增益
@app.route('/selectVolumeAdd',methods = ['POST'])
def selectVolumeAdd():
    global task_config

    scene_idx = request.form.get('scene_idx')
    volume_add = request.form.get('volume_add')

    task_config[scene_idx]['volume_add'] = volume_add

    return jsonify({'success':True,'message':''})


# 同步配置信息
@app.route('/update_task_config',methods = ['POST'])
def update_task_config():
    global task_config

    return jsonify(task_config)


# 查询音频库存
@app.route('/getAudioCount',methods = ['POST'])
def getAudioCount():
    global task_config,reply_sents_map

    scene_idx = request.form.get('scene_idx')

    if not task_config[scene_idx]['current_product']:
        return jsonify({'success':True,'count': -1})

    if not task_config[scene_idx]['live_platform']:
        return jsonify({'success':True,'count': -1})

    if not task_config[scene_idx]['current_streamer']:
        return jsonify({'success':True,'count': -1})
    
    if not os.path.exists(os.path.join('config','话术音频',task_config[scene_idx]['live_platform'],task_config[scene_idx]['current_product'],task_config[scene_idx]['current_streamer'],'主播')):
        return jsonify({'success':True,'count': -1})
    
    data = [file for file in os.listdir(os.path.join('config','话术音频',task_config[scene_idx]['live_platform'],task_config[scene_idx]['current_product'],task_config[scene_idx]['current_streamer'],'主播'))]

    return jsonify({'success':True,'count':len(data)})


# 返回主音频
main_audio_hist = []

def del_file(file_path):
    time.sleep(1)

    try:
        os.remove(file_path)
    except:
        time.sleep(5)
        os.remove(file_path)


@app.route('/download_main_audio', methods = ['POST'])
def download_main_audio():
    global task_config

    live_platform = request.form.get('platform')
    current_product = request.form.get('product')
    current_streamer = request.form.get('streamer')
    
    print(live_platform,current_product,current_streamer)
    
    if not live_platform or not current_product or not current_streamer:
        return '',404

    print(live_platform,current_product,current_streamer,'请求获取主音频')

    if not os.path.exists(os.path.join('config','话术音频',live_platform)):
        return '',404

    if not os.path.exists(os.path.join('config','话术音频',live_platform,current_product)):
        return '',404

    if not os.path.exists(os.path.join('config','话术音频',live_platform,current_product,current_streamer)):
        return '',404

    if not os.path.exists(os.path.join('config','话术音频',live_platform,current_product,current_streamer,'主播')):
        return '',404

    files = sorted([file for file in os.listdir(os.path.join('config','话术音频',live_platform,current_product,current_streamer,'主播')) if file.endswith('.wav') and not file in main_audio_hist])

    if not files:
        return '',404

    file_path = os.path.join('config','话术音频',live_platform,current_product,current_streamer,'主播',files[0])

    response =  send_file(file_path,as_attachment = True)

    main_audio_hist.append(files[0])

    task_thread = threading.Thread(target = del_file,args = (file_path,))
    task_thread.start()

    return response


if __name__ == '__main__':
    app.run(debug = True,port = 6868,host = '0.0.0.0')
